import logging
from contextlib import contextmanager

import redis

from app import create_app

# 设置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class RedisTokenPool:
    def __init__(self, tokens_key, tokens):
        self.tokens_key = tokens_key
        try:
            app = create_app()
            self.redis = redis.StrictRedis.from_url(app.config['REDIS_URL'])
            if not self.redis.exists(tokens_key):
                logger.info("Initializing new token pool in Redis")
                for token in tokens:
                    self.redis.lpush(tokens_key, token)
                logger.info(f"Added {len(tokens)} tokens to the pool")
        except redis.exceptions.ConnectionError as e:
            logger.error(f"Redis connection error during initialization: {e}")
            raise
        except Exception as e:
            logger.error(f"Unexpected error during RedisTokenPool initialization: {e}")
            raise

    @contextmanager
    def get_token(self, timeout=5):
        try:
            logger.info("Attempting to acquire token")
            token = self.redis.brpop(self.tokens_key, timeout)
            if token:
                return token[1]  # brpop 返回的是一个元组(key, value)
            else:
                raise TimeoutError("Failed to acquire token")
        except redis.exceptions.RedisError as e:
            logger.error(f"Redis error during token acquisition: {e}")
            raise
        except Exception as e:
            logger.error(f"Unexpected error during token acquisition: {e}")
            raise

    def release_token(self, token):
        """释放令牌回令牌池"""
        try:
            self.redis.lpush(self.tokens_key, token)
            logger.info(f"Token returned to the pool: {token}")
        except redis.exceptions.RedisError as e:
            logger.error(f"Redis error during token release: {e}")
        except Exception as e:
            logger.error(f"Unexpected error during token release: {e}")
